feat(kernel): thread Azure Entra OAuth (U2M + SP M2M) through the auth bridge - #919
feat(kernel): thread Azure Entra OAuth (U2M + SP M2M) through the auth bridge#919eric-wang-1990 wants to merge 4 commits into
Conversation
…h bridge The kernel auth bridge rejected azure-oauth and had no azure-sp-m2m path. Route both Azure auth types onto the kernel's generic OAuth flows (the kernel needs no Azure-specific code; PR databricks/databricks-sql-kernel#263 added the token_url/scope override plumbing this relies on): - azure-oauth (Azure AD U2M) -> oauth-u2m with the Azure app client id (96eecda7-...), redirect port 8030, and the {app_id}/user_impersonation offline_access delegated scope (via AzureOAuthEndpointCollection, honoring DATABRICKS_AZURE_TENANT_ID). The kernel discovers endpoints via the workspace /oidc redirector. (PECOBLR-4120) - azure-sp-m2m (Azure service principal) -> oauth-m2m with the Azure creds, an Entra v2.0 token_url, and the {effective_app_id}/.default scope. Requires an explicit azure_tenant_id (the kernel path does not auto-discover it). The management-token header / azure_workspace_resource_id are not applied on the kernel path -- no SQL connector uses them, matching Go and Node. (PECOBLR-4141) kernel_auth_kwargs now takes hostname (for the effective Azure app id); the client passes self._server_hostname. TDD: replaced the azure-oauth NotSupportedError test with routing tests and added a TestKernelAzureSpM2M suite (routing, required tenant/creds, federation client id). 50 bridge tests pass; black clean. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
There was a problem hiding this comment.
Pull request overview
This PR extends the use_kernel=True auth bridge to support Azure Entra (Azure AD) OAuth by routing the connector’s Azure auth types (azure-oauth U2M and azure-sp-m2m M2M) onto the kernel’s generic OAuth flows with Azure-specific overrides.
Changes:
- Route
auth_type="azure-oauth"to kerneloauth-u2musing the Azure OAuth app bundle and Azure scope mapping. - Route
auth_type="azure-sp-m2m"to kerneloauth-m2musing Entra v2.0 token URL and{effective_app_id}/.defaultscope; requireazure_tenant_id. - Thread
hostnameintokernel_auth_kwargs(viaKernelDatabricksClient.open_session) and add/adjust unit tests + changelog entry.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/databricks/sql/backend/kernel/auth_bridge.py |
Adds Azure auth-type routing and related override plumbing for kernel session kwargs. |
src/databricks/sql/backend/kernel/client.py |
Passes the server hostname into kernel_auth_kwargs during session open. |
tests/unit/test_kernel_auth_bridge.py |
Updates U2M tests for azure-oauth routing and adds azure-sp-m2m routing/validation tests. |
CHANGELOG.md |
Documents Azure Entra OAuth support on the kernel backend. |
Suppressed comments (1)
src/databricks/sql/backend/kernel/auth_bridge.py:262
auth_type='azure-sp-m2m'usesazure_client_id/azure_client_secret, but because this branch returns before the generic ambiguity guards, a caller can also supplyoauth_client_id/oauth_client_secretand have them silently ignored. This makes misconfigurations hard to diagnose.
if auth_type == "azure-sp-m2m":
azure_client_id = opts.get("azure_client_id")
azure_client_secret = opts.get("azure_client_secret")
azure_tenant_id = opts.get("azure_tenant_id")
if not (azure_client_id and azure_client_secret):
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if auth_type == "azure-oauth": | ||
| raise NotSupportedError( | ||
| "use_kernel=True does not support auth_type='azure-oauth' (Azure " | ||
| "AD U2M) yet: the kernel resolves OAuth endpoints only from the " | ||
| "workspace-native OIDC configuration and cannot drive the Azure AD " | ||
| "authorization/token flow. Use the Thrift backend (default) for " | ||
| "azure-oauth. Tracked by PECOBLR-4120." | ||
| redirect_port = opts.get("oauth_redirect_port") | ||
| caller_scopes = _normalize_scopes(opts.get("oauth_scopes")) |
| app_id = get_effective_azure_login_app_id(hostname or "") | ||
| kwargs = { |
There was a problem hiding this comment.
Verdict: 1 Medium
Routing logic is clean and well-tested overall. One medium concern: the new Azure branches return before the ambiguity guards, so azure-oauth + oauth_client_secret (or + credentials_provider) is silently accepted as U2M whereas the parallel databricks-oauth case raises NotSupportedError. Nit (not filed inline): the step-3 comment # Only databricks-oauth reaches here (azure-oauth rejected up front) in the U2M block is now stale — azure-oauth is routed up front, not rejected.
| # before the generic M2M/PAT/U2M routing below (azure-sp-m2m carries its | ||
| # creds in azure_* kwargs, not oauth_client_id/secret, so it would | ||
| # otherwise fall through to the final "unsupported" error). | ||
|
|
There was a problem hiding this comment.
🟡 Medium — The azure-oauth (and azure-sp-m2m) branches return before the ambiguity guards run, so conflicting signals that are rejected for the non-Azure U2M path are silently accepted for Azure.
Concretely, auth_type="azure-oauth" + oauth_client_secret reaches this branch, builds oauth-u2m kwargs, and returns — the secret is silently discarded and the browser flow runs. The exactly-parallel non-Azure case (auth_type="databricks-oauth" + oauth_client_secret) instead raises NotSupportedError("Ambiguous auth ...") a few lines below. The same gap lets azure-oauth + a custom credentials_provider bypass the credentials_provider-ambiguity guard. This is also a behavior regression: the prior code rejected any azure-oauth request up front, so a mis-specified azure-oauth + secret used to fail loudly and now succeeds with a different principal/flow than the caller likely intended.
Consider applying the same ambiguity checks (secret-with-U2M, credentials_provider-with-secret) to the Azure U2M branch before returning, so an ambiguous Azure request fails at session-open rather than silently choosing the browser flow.
The auth table marked the azure_* fields as Kernel-unsupported and claimed azure-oauth 'still works on the kernel' (it was actually rejected). Reflect the new routing: azure-sp-m2m + azure-oauth now work on the kernel path; azure_tenant_id is required there; the management token / azure_workspace_resource_id are not applied (matching Go/Node). Co-authored-by: Isaac Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean, well-documented routing with strong test coverage for both new Azure flows (routing, required tenant/creds, federation client id). One low-severity consistency note: the azure-oauth branch returns before the ambiguity guards, so oauth_client_secret / credentials_provider are silently ignored there, unlike the databricks-oauth U2M path which rejects them.
| # workspace /oidc redirector (which an Azure workspace redirects to Entra). | ||
| # The AAD delegated scope ({app_id}/user_impersonation [+ offline_access]) | ||
| # is synthesised via AzureOAuthEndpointCollection, which also honors the | ||
| # DATABRICKS_AZURE_TENANT_ID app-id override. PECOBLR-4120. |
There was a problem hiding this comment.
🔵 Low — The azure-oauth branch returns before the section-0 ambiguity guards, so it silently ignores auth signals that the databricks-oauth U2M path rejects loudly. Specifically:
auth_type="azure-oauth"+oauth_client_secret→ routes to U2M and drops the secret, whereasauth_type="databricks-oauth"+oauth_client_secretraisesNotSupportedError("Ambiguous auth …").auth_type="azure-oauth"+credentials_provider→ also silently ignored.
These combinations are nonsensical (Azure M2M is expressed via azure-sp-m2m + azure_client_secret, not oauth_client_secret), so the practical blast radius is small. But for consistency with the U2M guard — and to avoid a confusing silent browser-flow launch when a caller clearly intended a non-interactive credential — consider applying the same ambiguity check to azure-oauth, or documenting the intentional divergence. Not a correctness bug, just an inconsistency in the fail-loud contract.
The kernel is the auth core now: for azure-oauth the bridge forwards only auth_type='azure-oauth' (+ optional client_id/redirect_port passthrough), and the kernel pins the workspace v2.0 authorize/token endpoints, the Azure app client id, port 8030, and the user_impersonation scope. Drops the connector-side endpoint/scope construction (and the AzureOAuthEndpointCollection / PYSQL_OAUTH_AZURE_* imports) from the kernel path. Live-verified end-to-end against an Azure workspace. azure-sp-m2m still routes to oauth-m2m here pending the kernel's dedicated azure-sp-m2m variant. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Medium · 2 Low
Overall solid — the azure-sp-m2m v2.0 .default scope + v2.0 token endpoint correctly mirror the Thrift path's v1.0 resource= form, and the empty-hostname fallback is safe. Main issue is a stale module docstring (F1, medium) that claims the azure-oauth bridge synthesises the client id/port/scope via AzureOAuthEndpointCollection when the code actually just forwards the azure-oauth selector and lets the kernel own resolution — the inline comment already contradicts it. Plus two low doc/consistency nits.
| kernel's *generic* OAuth flows with Azure values as overrides (the | ||
| kernel needs no Azure-specific code): | ||
|
|
||
| - ``azure-oauth`` (U2M) → ``oauth-u2m`` with the Azure app client id |
There was a problem hiding this comment.
🟡 Medium — The module docstring here describes azure-oauth as routing → ``oauth-u2m`` with the Azure app client id (``96eecda7-…``), redirect port ``8030``, and the AAD delegated scope ``{app_id}/user_impersonation offline_access`` (via ``AzureOAuthEndpointCollection``…). That does not match what the code actually does. The azure-oauth branch (line 226) forwards auth_type="azure-oauth" (not oauth-u2m) and constructs none of those values — the inline comment at lines 218–225 explicitly states "this binding does NOT construct endpoints/scopes — it just passes auth_type='azure-oauth' plus any optional client_id / redirect_port passthrough" and that the kernel pins the client id / port 8030 / scope.
So the top-of-file docstring directly contradicts both the executable code and the inline comment beside it: it claims the bridge maps to oauth-u2m and synthesises the Azure bundle, when in fact it delegates all of that to the kernel via the azure-oauth selector. A maintainer reading the module docstring will expect AzureOAuthEndpointCollection.get_scopes_mapping to be called here (it never is on this path). Please align the docstring with the actual selector-forwarding behavior.
| # a malformed Entra URL), the Azure app client id, port 8030, and the | ||
| # `{app_id}/user_impersonation offline_access` scope. So this binding does | ||
| # NOT construct endpoints/scopes — it just passes `auth_type='azure-oauth'` | ||
| # plus any optional client_id / redirect_port passthrough. PECOBLR-4120. |
There was a problem hiding this comment.
🔵 Low — Asymmetric handling of auth_type + oauth_client_secret. For databricks-oauth, supplying oauth_client_secret (a U2M selector plus an M2M secret) is deliberately rejected as "Ambiguous auth" so it doesn't silently authenticate against the wrong principal. But the azure-oauth branch runs before those guards and simply ignores any oauth_client_secret present, silently proceeding with the U2M browser flow. Given azure-oauth is newly supported here, the same class of conflicting-signal request should arguably fail loudly for consistency rather than be silently dropped.
(Anchored to the nearest changed line — see the description for the exact location.)
| kernel can't drive the Azure AD authorization/token flow — so it is | ||
| rejected up front (PECOBLR-4120).""" | ||
| ``azure-oauth`` (Azure AD U2M) routes here too — see | ||
| ``test_azure_oauth_routes_to_kernel_u2m`` (PECOBLR-4120).""" |
There was a problem hiding this comment.
🔵 Low — The TestKernelOAuthU2M class docstring points readers to test_azure_oauth_routes_to_kernel_u2m, but no such test exists. The azure-oauth tests are actually named test_azure_oauth_forwards_selector_kernel_owns_resolution and test_azure_oauth_honors_custom_client_id_and_port. Update the cross-reference so it names a real test.
…optional mgmt token
Make the `azure-sp-m2m` bridge thin, matching the kernel becoming the
Azure-aware auth core. The connector now forwards
`auth_type='azure-sp-m2m'` + `azure_client_id` / `azure_client_secret`
(and optional `azure_tenant_id` / `azure_workspace_resource_id`) straight
to the kernel Session, instead of constructing the Entra token endpoint
and `{app_id}/.default` scope itself.
Behavior changes on the kernel path (Thrift parity):
- `azure_tenant_id` is now OPTIONAL — the kernel auto-discovers the
tenant from the workspace's `/aad/auth` redirect when omitted, exactly
as the Thrift backend does. (Previously the kernel path required it.)
- `azure_workspace_resource_id` is now honored as an optional add-on:
forward it and the kernel fetches an Azure-management token and sends
the `X-Databricks-Azure-SP-Management-Token` +
`X-Databricks-Azure-Workspace-Resource-Id` pair, so an SP with only an
Azure RBAC role (not a workspace member) can authenticate. (Previously
it was dropped with a warning.)
Also thread the `azure_*` connection kwargs into `kernel_auth_options`
in session.py — without this the bridge never received them and
`azure-sp-m2m` failed at session-open with "requires azure_client_id".
Adds a regression test for that threading, and rewrites the bridge tests
for thin forwarding (tenant optional, resource id forwarded). Drops the
now-unused `get_effective_azure_login_app_id` import and
`_AZURE_AAD_LOGIN_HOST` constant.
Co-authored-by: Isaac
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Medium · 1 Low
Looks good overall — the Azure routing logic is thin, well-tested, and correct. Two cleanup items: an unused hostname parameter now threaded through client.py for no effect (medium), and a dead test import (low). Nit: the comment at auth_bridge.py:318 ("azure-oauth rejected up front") is now stale — azure-oauth is routed up front, not rejected.
| def kernel_auth_kwargs( | ||
| auth_provider: Optional[AuthProvider], | ||
| auth_options: Optional[Dict[str, Any]] = None, | ||
| hostname: Optional[str] = None, |
There was a problem hiding this comment.
🟡 Medium — The new hostname parameter is never used. kernel_auth_kwargs now accepts hostname (line 152) and client.py was changed to pass self._server_hostname, but the function body never references hostname on any code path — the azure-oauth and azure-sp-m2m branches forward the selector/creds verbatim and let the kernel own resolution.
The PR description says the param exists "for the effective Azure app id per cloud," which matches the intent behind the test-file import of get_effective_azure_login_app_id — but that computation was ultimately delegated to the kernel, leaving this parameter (and the client.py threading) as dead plumbing. Either drop the parameter and the client.py call-site change, or use it. Threading a value in that the function ignores invites future readers to assume it affects resolution when it does not.
| PYSQL_OAUTH_SCOPES, | ||
| PYSQL_OAUTH_REDIRECT_PORT_RANGE, | ||
| ) | ||
| from databricks.sql.auth.common import get_effective_azure_login_app_id |
There was a problem hiding this comment.
🔵 Low — get_effective_azure_login_app_id is imported here but never referenced anywhere in the test module (the TestKernelAzureSpM2M suite asserts only on the thin-forwarding kwargs, which don't involve the effective app id). This is an unused import — flake8/linters will flag it, and it's a leftover from an earlier design where the bridge computed the app id itself. Remove the import.
What
Teaches the
use_kernel=Trueauth bridge to support Azure Entra (Azure AD) OAuth — both the U2M browser flow (azure-oauth) and the service-principal M2M flow (azure-sp-m2m). Previouslyazure-oauthwas rejected withNotSupportedErrorandazure-sp-m2mfell through to the generic "unsupported" error.The kernel needs no Azure-specific code: both types route onto the kernel's generic OAuth flows with Azure values supplied as overrides. This depends on the kernel-side override plumbing added in databricks/databricks-sql-kernel#263 (
token_url/ scope overrides on M2M + U2M; M2M skips OIDC discovery when a token endpoint is pinned).How
src/databricks/sql/backend/kernel/auth_bridge.py:azure-oauth(U2M) → kerneloauth-u2mwith the Azure app client id (96eecda7-…), redirect port8030, and the AAD delegated scope{app_id}/user_impersonation offline_access(synthesised viaAzureOAuthEndpointCollection.get_scopes_mapping, honoringDATABRICKS_AZURE_TENANT_ID). The kernel runs the browser flow and discovers endpoints via the workspace/oidcredirector (which an Azure workspace redirects to Entra). PECOBLR-4120.azure-sp-m2m(M2M) → kerneloauth-m2mwith the Azure service-principal credentials, an Entra v2.0token_url(login.microsoftonline.com/{tenant}/oauth2/v2.0/token), and the{effective_app_id}/.defaultscope. PECOBLR-4141.azure_tenant_id(the kernel path does not auto-discover the tenant from the workspace's/aad/authredirect as the Thrift path does) — raises a clear error otherwise.X-Databricks-Azure-SP-Management-Token) andazure_workspace_resource_idare not applied on the kernel path. No SQL connector (Go, Node) uses them; the Databricks-audience token authenticates SPs that are workspace principals (the SQL norm). An RBAC-only-not-in-workspace SP is unsupported here (documented; warns ifazure_workspace_resource_idis set).kernel_auth_kwargsnow takeshostname(for the effective Azure app id per cloud);client.pypassesself._server_hostname.Tests
TDD — replaced the
azure-oauthNotSupportedErrortest with routing tests, and added aTestKernelAzureSpM2Msuite (routing, required tenant, required creds, federation client id). 50 bridge tests pass;blackclean. (test_kernel_client.pyskips without the kernel wheel, as before.)Related
This pull request and its description were written by Isaac.
This PR was created with GitHub MCP.